#!/usr/bin/python3
# coding=utf-8

import sys
import akshare as ak
import pandas as pd
import numpy as np
import psycopg2
import talib
import time
import datetime
import matplotlib.pyplot as plt

SD=20
WIDTH=2

# DIS:=STD(CLOSE,SD);
# MID:MA(CLOSE,SD),COLORFFAEC9;
# UPPER:MID+WIDTH*DIS,COLORFFC90E;
# LOWER:MID-WIDTH*DIS,COLOR0CAEE6;


conn = psycopg2.connect(database="finance",
                        user="postgres", password="123456",
                        host="127.0.0.1", port="5432")
cur = conn.cursor()  # 创建指针对象

stock_rsi = []
sql = ""

cur.execute("SELECT code, sname FROM stock_market.stock")
datas = cur.fetchall()
for item in datas:
    code = item[0]
    sname = item[1]
    if "ST" in sname:
        print("不计算 %s", sname)
        continue
    if "退" in sname:
        print("不计算 %s", sname)
        continue
    try:
        select_sql = '''SELECT date,open,close,hight,low FROM stock_datas.stock_qfq_{} ORDER BY date DESC LIMIT {};'''
        sql = select_sql.format("000333", 20000)
        # print(sql)
        cur.execute(sql)
        df = pd.DataFrame(cur.fetchall())
        if df.values.size <= 10:
            print("======", code, sname, df.values.size)
            continue
        df.columns= ["date", "open", "close", "high", "low"]
        new_df = df.sort_values("date", ascending=True)
        new_df = new_df.rename(index=new_df["date"][:])
        col = new_df.iloc[:, 2]
        arr = col.values
        upperband, middleband, lowerband = talib.BBANDS(arr, timeperiod=20, nbdevup=2.0, nbdevdn=2.0, matype=0)
        up = np.array(upperband)
        low = np.array(lowerband)
        cha = up-low
        new_df["cha"] = cha
        print(upperband)
        print(lowerband)
        query_index = pd.date_range("20220901", "20220930", freq='D')
        days_index = []
        for day in query_index:
            if day in new_df.index.values:
                days_index.append(day.date())
        plt.plot(new_df.loc[days_index, 'cha'])
        plt.show()
        sys.exit()
        # rsi = talib.RSI(arr, timeperiod=14)
        # if len(rsi) < 100:
        #     continue
        # rsi20 = rsi[-3:]
        # for item in rsi20:
        #     if item <= 20.0:
        #         stock_rsi.append(code)
        #         print("insert", code, rsi20)
        #         break
        # new_df.insert(6, "rsi", rsi)
        # print(new_df.tail(30))
    except Exception as e:
        print(sql)
        cur.execute("ROLLBACK")
        print('errmsg： {}'.format(e))
        sys.exit(0)
    # time.sleep(0.001)
print("============================================")
print(stock_rsi)
print("============================================")
